//
// See the file COPYING for redistribution information.
//

#include <config.h>
#include <cstring>
#include <cerrno>
#include <netdb.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <arpa/nameser.h>
#include <arpa/nameser_compat.h>
#include <ares_dns.h>
#include <Naipax/Common.hpp>
#include <Naipax/Error.hpp>
#include <Naipax/Ares.hpp>

using namespace std;
using namespace Naipax;

unsigned char* _ParseQuestion(const unsigned char *pBufPtr,
                              const unsigned char *pBuf, int32_t BufLen) {
    char *pName = NULL;
    long Len = 0;
    if (ares_expand_name(pBufPtr, pBuf, BufLen,
                         &pName, &Len) != ARES_SUCCESS) {
        return NULL;
    }
    ares_free_string(pName);
    pBufPtr += Len;
    if (pBufPtr + QFIXEDSZ > pBuf + BufLen) {
        return NULL;
    }
    pBufPtr += QFIXEDSZ;
    return (unsigned char*) pBufPtr;
}

unsigned char* _ParseRR(const unsigned char *pBufPtr,
                        const unsigned char *pBuf, int32_t BufLen,
                        int32_t &rTtl) {
    char *pName = NULL;
    long Len = 0;
    if (ares_expand_name(pBufPtr, pBuf, BufLen,
                         &pName, &Len) != ARES_SUCCESS) {
        return NULL;
    }
    ares_free_string(pName);
    pBufPtr += Len;
    if (pBufPtr + RRFIXEDSZ > pBuf + BufLen) {
        return NULL;
    }
    rTtl = DNS_RR_TTL(pBufPtr);
    pBufPtr += RRFIXEDSZ;
    if (pBufPtr + DNS_RR_LEN(pBufPtr) > pBuf + BufLen) {
        return NULL;
    }
    return (unsigned char*) pBufPtr + DNS_RR_LEN(pBufPtr);
}

void _QueryCallback(void *pArg, int32_t Status, 
                    unsigned char *pBuf, int32_t BufLen) {
    if (pArg == NULL) {
        NAIPAX_THROW(Error::INVALID_ARGUMENT);
    }
    AresData *pDnData = (AresData*) pArg;

    if (Status != ARES_SUCCESS) {
        pDnData->mAresReturn = ~ARES_SUCCESS;
        pDnData->mIsDone = true;
        return;
    }

    // Won't happen. For security. Please see also adig.c in the
    // c-ares library.
    if (BufLen < HFIXEDSZ) {
        pDnData->mAresReturn = ~ARES_SUCCESS;
        pDnData->mIsDone = true;
        return;
    }
    
    unsigned char *pBufPtr = pBuf;
    pBufPtr = pBuf + HFIXEDSZ;
    pBufPtr = _ParseQuestion(pBufPtr, pBuf, BufLen);
    if (pBufPtr != NULL && DNS_HEADER_ANCOUNT(pBuf) > 0) {
        int32_t Ttl;
        pBufPtr = _ParseRR(pBufPtr, pBuf, BufLen, Ttl);
        if (pBufPtr == NULL) {
            pDnData->mAresReturn = ~ARES_SUCCESS;
            pDnData->mIsDone = true;
            return;
        }
        pDnData->mTtl = Ttl;
    }

    struct hostent *HostInfo = NULL;
    int32_t Ret = ares_parse_a_reply(pBuf, BufLen, &HostInfo);
    if (Ret == ARES_SUCCESS) {
        if (HostInfo->h_addrtype == AF_INET) {
            pDnData->mIsIpV6 = false;
            pDnData->mIsDone = true;
            pDnData->mAresReturn = ARES_SUCCESS;
            memcpy(&(pDnData->mIpV4Addr), (HostInfo->h_addr_list[0]),
                   sizeof(struct in_addr));
        }
        else {
            pDnData->mIsIpV6 = true;
        }
    }
    else {
        pDnData->mIsValid = false;
        pDnData->mIsDone = false;
    }
    if (HostInfo != NULL) {
        ares_free_hostent(HostInfo);
    }
}

AresData::AresData(const char *pHostname, size_t HostnameLength)
        : mIsValid(false), mIsDone(false) {
    Build(pHostname, HostnameLength);
}

AresData::AresData() : mIsValid(false), mIsDone(false) {
}

void AresData::Build(const char *pHostname, size_t HostnameLength) {
    if (pHostname != NULL) {
        if (HostnameLength < HOSTNAME_SIZE) {
            strncpy(mHostname, pHostname, HostnameLength);
            mHostname[HostnameLength] = '\0';
            mIsValid = true;
        }
        else {
            NAIPAX_THROW(Error::DATA_OUT_OF_RANGE);
        }
    }
    else {
        NAIPAX_THROW(Error::INVALID_ARGUMENT);
    }
}

Ares::Ares(unsigned int ConnCount) : _mConnCount(ConnCount) {
    try {
        _mpAresChannels = new ares_channel[ConnCount];
        _mpDataArray = new AresData[ConnCount];
        _mpAresDataCache = new Cache(Naipax::HOSTNAME_SIZE, sizeof(AresData),
                                     104729);
    }
    catch (std::bad_alloc &E) {
        NAIPAX_THROW(Error::CANNOT_INIT_ARES);
    }
    for (uint32_t i = 0; i < ConnCount; ++i) {
        struct ares_options Options;
        Options.timeout = 3;
        Options.flags = ARES_FLAG_STAYOPEN;
        if (ares_init_options(&(_mpAresChannels[i]), &Options,
                              ARES_OPT_TIMEOUT
                              | ARES_OPT_FLAGS) != ARES_SUCCESS) {
            NAIPAX_THROW(Error::CANNOT_INIT_ARES);
        }
    }
}

Ares::~Ares() {
    delete _mpAresDataCache;
    delete[] _mpDataArray;
    for (uint32_t i = 0; i < _mConnCount; ++i) {
        ares_destroy(_mpAresChannels[i]);
    }
    delete[] _mpAresChannels;
}

int32_t Ares::Submit(const AresData &rData) {
    size_t HostnameLen = strlen(rData.mHostname);

    if (HostnameLen == 0) {
        NAIPAX_THROW(Error::INVALID_ARGUMENT);
    }
    
    // Try to load data from cache
    AresData CachedAresData;
    uint32_t CachedAresDataSize = 0;

    if (_mpAresDataCache->Get(rData.mHostname, HostnameLen,
                              &CachedAresData, &CachedAresDataSize) == true) {
        _mAresDataVector.push_back(CachedAresData);
        return 0;
    }
    
    uint32_t i = 0;
    for (i = 0; i < _mConnCount; ++i) {
        if (_mpDataArray[i].mIsValid == true
            && time(NULL) > _mpDataArray[i].mExpireTime) {
            _mpDataArray[i].mIsValid = false;
        }
        if (_mpDataArray[i].mIsValid == false) {
            break;
        }
    }
    // If there is no empty slot.
    if (i == _mConnCount) {
        NAIPAX_THROW(Error::BUFFER_IS_FULL);
    }
    _mpDataArray[i] = rData;
    _mpDataArray[i].mIsValid = true;
    _mpDataArray[i].mIsDone = false;
    _mpDataArray[i].mExpireTime = time(NULL) + 3;
    
    ares_query(_mpAresChannels[i], DataArray()[i].mHostname, C_IN, T_A,
               _QueryCallback, &(DataArray()[i]));
    return 0;
}

bool Ares::HasUnresolved() {
    for (uint32_t i = 0; i < _mConnCount; ++i) {
        if (_mpDataArray[i].mIsValid == true) {
            return true;
        }
    }
    return false;
}

bool Ares::HasResult(vector<AresData> &rDataVector) {
    // Check _mAresDataVector
    if (_mAresDataVector.size() > 0) {
        rDataVector = _mAresDataVector;
        _mAresDataVector.clear();
    }

    for (uint32_t i = 0; i < _mConnCount; ++i) {
        if (_mpDataArray[i].mIsValid != true
            || (_mpDataArray[i].mIsValid == true
                && time(NULL) > _mpDataArray[i].mExpireTime)) {
            _mpDataArray[i].mIsValid = false;
            continue;
        }

        FD_ZERO(&_mReadFdSet);
        FD_ZERO(&_mWriteFdSet);
        int32_t NFds = ares_fds(_mpAresChannels[i],
                                &_mReadFdSet, &_mWriteFdSet);
        if (NFds == 0) {
            continue;
        }
        
        struct timeval Timeout;
        Timeout.tv_sec = 0;
        Timeout.tv_usec = 100;

        int32_t Count = select(NFds, 
                               &_mReadFdSet, &_mWriteFdSet, NULL, &Timeout);
        
        if (Count <= 0) {
            continue;
        }
        ares_process(_mpAresChannels[i], &_mReadFdSet, &_mWriteFdSet);

        if (_mpDataArray[i].mIsValid == true
            && _mpDataArray[i].mIsDone == true) {
            if (_mpDataArray[i].mAresReturn == ARES_SUCCESS) {
                rDataVector.push_back(_mpDataArray[i]);
                _mpAresDataCache->Set(_mpDataArray[i].mHostname,
                                      strlen(_mpDataArray[i].mHostname),
                                      &(_mpDataArray[i]), sizeof(AresData),
                                      86400);
            }
            _mpDataArray[i].mIsValid = false;
            _mpDataArray[i].mIsDone = false;
        }
    }
    
    if (rDataVector.size() == 0) {
        return false;
    }
    else {
        return true;
    }
}

void Ares::Cancel() {
    // The type of the return value of ares_cancel() is int, but the
    // values are not documented in the header file. Ignore it for
    // now.
    for (uint32_t i = 0; i < _mConnCount; ++i) {
        ares_cancel(_mpAresChannels[i]);
    }
}
